from os import path
from collections import defaultdict
import numpy as np
from ppgnss import gnss_utils
import matplotlib.pyplot as plt

obj_dir = "../data/snr"
sites = ["B216", "B240", "B286", "PD02"]
doys = range(98, 105)
for site in sites:
    for idoy, doy in enumerate(doys):
        obj_filename = path.join(obj_dir, "%s%03d.obj" % (site, doy))
        data = gnss_utils.loadobject(obj_filename)
        xr_ele, xr_azi, xr_snr = data["ele"], data["azi"], data["snr"]
        ele_min = 1.25
        azi_min = 2.5
        ele_step_size = 2.5
        azi_step_size = 5
        nrows = int(90/ele_step_size)
        ncols = int(360/azi_step_size)
        nsat = len(xr_ele.coords["prn"])
        invalid_idx = xr_ele < 0
        idx_row_grid = np.array(np.int32(np.floor(xr_ele/ele_step_size)))
        idx_col_grid = np.array(np.int32(np.floor(xr_azi/azi_step_size))) 
        xx_grid = azi_min + idx_col_grid * azi_step_size
        yy_grid = ele_min + idx_row_grid * ele_step_size
        eles = np.arange(ele_min, 90, ele_step_size)
        azis = np.arange(azi_min, 360, azi_step_size)
        eles = np.linspace(ele_min, 90-ele_min, nrows)
        azis = np.linspace(azi_min, 360-azi_min, ncols)
        azis_grid, eles_grid = np.meshgrid(azis, eles)
        idx_row_grid[invalid_idx] = -9999
        idx_col_grid[invalid_idx] = -9999
        fig = plt.figure(figsize=(80/2.54, 80/2.54))
        fig, axes = plt.subplots(nrows=3, ncols=3, subplot_kw={'projection': 'polar'})
        for iband, band in enumerate(xr_snr.coords["band"].values):
            strband = str(band)
            sattype = strband[:3]
            idx_gnss = xr_ele.coords["gnss"].values == sattype
            band_grid = np.zeros((nrows, ncols))
            count_grid = np.zeros((nrows, ncols))
            for iepoch, epoch_snr in enumerate(xr_snr.loc[strband]):
                iirow = idx_row_grid[idx_gnss, iepoch].reshape(-1)
                iicol = idx_col_grid[idx_gnss, iepoch].reshape(-1)
                satlist = np.array(range(nsat))
                valid_idx_sat = iirow > 0
                for isat, irow, icol in zip(satlist[valid_idx_sat], iirow[valid_idx_sat], iicol[valid_idx_sat]):
                    band_grid[irow][icol] += xr_snr.loc[band][iepoch][isat].values
                    count_grid[irow][icol] += 1
            mean_grid = band_grid/count_grid
            mean_grid[count_grid==0] = np.nan
            
            i = iband//3
            j = iband%3
            axes[i][j].pcolor(np.deg2rad(azis_grid), np.pi/2-np.deg2rad(eles_grid), mean_grid) 
            axes[i][j].set_theta_zero_location("N")
            axes[i][j].set_theta_direction(-1)

        fig_filename = "%s_%03d_mean.png" % (site, doy)
        plt.savefig(fig_filename, dpi=300)
        print(fig_filename + " is saved!")
        plt.close()